import sys

sys.path.append("../../../")

from common.python.common import *
from common.python.ipe_param import *


class IpeParamScrfd(IpeParamBase):
    def __init__(self, model_width, model_height):
        super().__init__(model_width, model_height)
        # 图片resize参数
        self.m_iResizeWidth = 0
        self.m_iResizeHeight = 0
        # 图片填充参数
        self.m_iPadRight = 0
        self.m_iPadBottom = 0

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_c2c_config(
                self.m_oConfigDesc, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0
            )
        )
        error_check(
            sdk.lyn_ipe_set_resize_config(
                self.m_oConfigDesc, self.m_iResizeWidth, self.m_iResizeHeight
            )
        )
        error_check(
            sdk.lyn_ipe_set_pad_config(
                self.m_oConfigDesc,
                0,
                self.m_iPadRight,
                self.m_iPadBottom,
                0,
                114,
                114,
                114,
            )
        )
        error_check(
            sdk.lyn_ipe_cal_output_pic_desc(
                self.m_oPicDescOut, self.m_oPicDescIn, self.m_oConfigDesc, 2
            )
        )
        error_check(
            sdk.lyn_ipe_process_async(
                stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
            )
        )

    def set_img_info(self, width, height, format):
        scale = max(width / self.m_iModelWidth, height / self.m_iModelHeight)
        self.m_iResizeWidth = int(width / scale) & ~1
        self.m_iResizeHeight = int(height / scale) & ~1

        if self.m_iResizeWidth > 1920:
            self.m_iResizeWidth = 1920
        if self.m_iResizeHeight > 1080:
            self.m_iResizeHeight = 1080

        self.m_iPadRight = self.m_iModelWidth - self.m_iResizeWidth
        self.m_iPadBottom = self.m_iModelHeight - self.m_iResizeHeight

        super().set_img_info(width, height, format)
